package com.cfp4cloud.cfp.knowledge.support.function;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import com.cfp4cloud.cfp.common.security.service.CfpUser;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.cfp4cloud.cfp.common.core.util.R;
import com.cfp4cloud.cfp.knowledge.dto.BaseAiRequest;
import com.cfp4cloud.cfp.knowledge.dto.ChatMessageDTO;
import com.cfp4cloud.cfp.knowledge.entity.AiDataEntity;
import com.cfp4cloud.cfp.knowledge.service.AiChartGenerateService;
import com.cfp4cloud.cfp.knowledge.service.AiDataService;
import com.cfp4cloud.cfp.knowledge.support.annotation.FieldPrompt;
import com.cfp4cloud.cfp.knowledge.support.feign.RemoteTableInfoService;
import com.cfp4cloud.cfp.knowledge.support.provider.ChatMemoryAdvisorProvider;
import com.cfp4cloud.cfp.knowledge.support.util.ChatMessageContextHolder;
import com.cfp4cloud.cfp.knowledge.support.util.PromptBuilder;
import dev.langchain4j.internal.Utils;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.Select;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;

/**
 * Chat2SQL函数调用实现类 该类实现了将自然语言转换为SQL语句并执行的功能 支持SQL查询结果的表格展示和图表生成
 *
 * @author chenda
 * @date 2024/07/30
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class Chat2SqlFunctionCalling implements FunctionCalling<SqlRequest> {

	private final ChatMemoryAdvisorProvider chatMemoryAdvisorProvider;

	private final AiChartGenerateService chartGenerateService;

	private final AiDataService aiDataService;

	/**
	 * 是否在前端展示此函数 返回false表示该函数不会在前端功能列表中显示
	 * @return boolean false-不展示
	 */
	@Override
	public boolean showFunction() {
		return false;
	}

	/**
	 * 检查参数是否正确 主要验证SQL语句的合法性和安全性 只允许执行SELECT查询语句，拒绝其他操作语句
	 * @param sqlRequest SQL请求对象，包含要执行的SQL语句
	 * @param userDetails 当前用户信息
	 * @param extDetails 扩展详情，包含数据源等信息
	 * @return R<String> 参数检查结果
	 */
	@Override
	public R<String> checkParams(SqlRequest sqlRequest, CfpUser userDetails, ChatMessageDTO.ExtDetails extDetails) {

		log.info("Chat2SQL 结果：{}", sqlRequest);

		Statement statement = null;
		try {
			statement = CCJSqlParserUtil.parse(sqlRequest.getSql());
		}
		catch (JSQLParserException e) {
			return R.failed(StrUtil.format("SQL: {} 解析失败", sqlRequest.getSql()));
		}
		if (statement instanceof Select) {
			sqlRequest.setSql(statement.toString());
			return R.ok();
		}

		return R.failed(StrUtil.format("SQL: {} 解析失败，安全问题不支持非查询语句", sqlRequest.getSql()));
	}

	/**
	 * 处理SQL执行请求 1. 执行SQL查询 2. 处理查询结果 3. 如果需要，生成图表 4. 将结果转换为Markdown表格格式返回
	 * @param sqlRequest SQL请求对象
	 * @param userDetails 用户详细信息
	 * @param extDetails 扩展详细信息，包含数据源ID等
	 * @return R<String> 执行结果，包含Markdown格式的表格数据
	 */
	@Override
	public R<String> handle(SqlRequest sqlRequest, CfpUser userDetails, ChatMessageDTO.ExtDetails extDetails) {
		RemoteTableInfoService remoteTableInfoService = SpringUtil.getBean(RemoteTableInfoService.class);
		AiDataEntity dataEntity = aiDataService.getById(extDetails.getDataId());
		R execSQLResult = remoteTableInfoService.execSQL(dataEntity.getDsName(), sqlRequest.getSql());

		// 处理异常情况
		if (!execSQLResult.isOk()) {
			// 异常情况下，清空多轮会话,避免后续丢失精度
			ChatMessageDTO chatMessageDTO = ChatMessageContextHolder.get();
			chatMemoryAdvisorProvider.get(chatMessageDTO.getConversationId()).clear();

			return R.failed(StrUtil.format("""
					执行后结果为空，不进行表格和图表渲染。
					```sql
					{}
					```
					请进行条件调整后重试，或联系管理员进行字段评估
					```bash
					{}
					```
					""", sqlRequest.getSql(), execSQLResult.getMsg()));
		}

		List<Map<String, Object>> resultData = (List<Map<String, Object>>) execSQLResult.getData();

		if (CollUtil.isEmpty(resultData)) {
			// 为空判定为：异常情况，清空多轮会话,避免后续丢失精度
			ChatMessageDTO chatMessageDTO = ChatMessageContextHolder.get();
			chatMemoryAdvisorProvider.get(chatMessageDTO.getConversationId()).clear();

			return R.failed(StrUtil.format("""
					执行后结果为空，不进行表格和图表渲染。
					```sql
					{}
					```
					请进行条件调整后重试，或联系管理员进行字段评估
					""", sqlRequest.getSql()));
		}

		// 如果建议图表输出,则执行异步的图表计算逻辑
		if (sqlRequest.isChartOutput() && sqlRequest.getChartType() > 0 && sqlRequest.getChartType() <= 3) {
			ChatMessageDTO chatMessageDTO = ChatMessageContextHolder.get();
			chatMessageDTO.getExtDetails().setChartType(String.valueOf(sqlRequest.getChartType()));
			chatMessageDTO.getExtDetails().setChartId(Utils.randomUUID());
			ChatMessageContextHolder.set(chatMessageDTO);
			chartGenerateService.generateChart(chatMessageDTO, execSQLResult.getData());
		}

		String markdownTable = PromptBuilder.toMarkdownTable(resultData);
		log.info("SQL 执行结果：{}", markdownTable);
		return R.ok(markdownTable);

	}

}

/**
 * SQL请求参数类 用于接收AI生成的SQL查询相关参数
 */
@Data
@JsonInclude(JsonInclude.Include.NON_NULL)
class SqlRequest extends BaseAiRequest {

	/**
	 * AI生成的可执行SQL语句
	 */
	@FieldPrompt("可以执行的SQL语句")
	private String sql;

	/**
	 * 是否建议使用图表形式展示查询结果
	 */
	@FieldPrompt(value = "是否推荐用图表的形式输出结果")
	private boolean chartOutput;

	/**
	 * 推荐的图表类型 1: 折线图(Line chart) 2: 饼图(Pie chart) 3: 柱状图(Bar chart)
	 */
	@FieldPrompt(value = "如果推荐图表输出的话，更适合哪种图表? 1. Line chart 2. Pie chart 3. Bar chart")
	private int chartType;

}
